import torch
import torch.nn as nn
from einops import rearrange

x_enc = torch.rand(3, 16, 2560)
mean = torch.mean(x_enc)  # 均值
std = torch.std(x_enc)
x_enc = (x_enc - mean) / std  # 标准差
B, L, enc_in = x_enc.shape

x_enc_r = rearrange(x_enc, 'b l (k o) -> b l k o', o=2)  # 将输入数据重排为四维张量，可能是为了将实部和虚部分开处理。
x_enc_complex = torch.complex(x_enc_r[:, :, :, 0], x_enc_r[:, :, :, 1])  # 将实部和虚部组合成复数张量
x_enc_delay = torch.fft.ifft(x_enc_complex, dim=2)  # 对复数信道进行逆傅立叶变换，得到延迟域的表示
x_enc_delay = torch.cat([torch.real(x_enc_delay), torch.imag(x_enc_delay)], dim=2)
print(x_enc_delay.shape)
x_enc_delay = x_enc_delay.reshape(B, L // 4, 4, enc_in) # B L/4 4 D
print(x_enc_delay.shape)
patch_layer = nn.Linear(4, 4)
x_enc_delay = x_enc_delay.permute(0, 1, 3, 2)
print(x_enc_delay.shape)
x_enc_delay = patch_layer(x_enc_delay).permute(0, 1, 3, 2)
print(x_enc_delay.shape)
x_enc_delay = x_enc_delay.reshape(B, L, enc_in)
print(x_enc_delay.shape)
x_enc_delay = rearrange(x_enc_delay, 'b l (k o) -> b o l k', o=2)
print(x_enc_delay.shape)
res_dim = 64
RB_f = nn.Sequential(nn.Conv2d(2, res_dim, 3, 1, 1))
x_enc_delay = RB_f(x_enc_delay)
print(x_enc_delay.shape)
class DepthwiseSeparableConv(nn.Module):
    def __init__(self, in_planes, out_planes):
        super(DepthwiseSeparableConv, self).__init__()
        self.depthwise = nn.Conv2d(in_planes, in_planes, kernel_size=3, padding=1, groups=in_planes)
        self.pointwise = nn.Conv2d(in_planes, out_planes, kernel_size=1)

    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

class Res_block(nn.Module):
    def __init__(self, in_planes, out_planes):
        super(Res_block, self).__init__()
        self.conv1 = DepthwiseSeparableConv(in_planes, out_planes)
        self.bn1 = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = DepthwiseSeparableConv(out_planes, out_planes)
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.ca = ChannelAttention(out_planes, ratio=1)

    def forward(self, x):
        rs1 = self.relu(self.bn1(self.conv1(x)))
        rs1 = self.bn2(self.conv2(rs1))
        channel_attn = self.ca(rs1)
        output = channel_attn * rs1
        rs = torch.add(x, output)
        return rs
class Res_block(nn.Module):
    def __init__(self, in_planes, out_planes, use_1x1conv=False):
        super(Res_block, self).__init__()
        self.conv1 = DepthwiseSeparableConv(in_planes, out_planes)
        self.bn1 = nn.BatchNorm2d(out_planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = DepthwiseSeparableConv(out_planes, out_planes)
        self.bn2 = nn.BatchNorm2d(out_planes)
        self.ca = ChannelAttention(out_planes, ratio=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(in_planes, out_planes, kernel_size=1)
        else:
            self.conv3 = None

    def forward(self, x):
        rs1 = self.relu(self.bn1(self.conv1(x)))
        rs1 = self.bn2(self.conv2(rs1))
        channel_attn = self.ca(rs1)
        output = channel_attn * rs1
        if self.conv3:
            x = self.conv3(x)
        rs = torch.add(x, output)
        return rs

# 其中的一个优化
class Res_block(nn.Module):
    def __init__(self, in_planes, use_1x1conv=False):
        super(Res_block, self).__init__()
        self.conv1 = DepthwiseSeparableConv(in_planes, in_planes)
        self.bn1 = nn.BatchNorm2d(in_planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = DepthwiseSeparableConv(in_planes, in_planes)
        self.bn2 = nn.BatchNorm2d(in_planes)
        self.ca = ChannelAttention(in_planes, ratio=1)
        if use_1x1conv:
            self.conv3 = nn.Conv2d(in_planes, in_planes, kernel_size=1)
        else:
            self.conv3 = None

    def forward(self, x):
        rs1 = self.relu(self.bn1(self.conv1(x)))
        rs1 = self.bn2(self.conv2(rs1))
        channel_attn = self.ca(rs1)
        output = channel_attn * rs1
        if self.conv3:
            x = self.conv3(x)
        rs = torch.add(x, output)
        return rs

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=4):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)

        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)